# Modified from https://pjreddie.com/media/files/voc_label.py

import os
from os import getcwd
# import xml.etree.ElementTree as ET

sets = [
    ("2012", "train"),
    ("2012", "val"),
    ("2007", "train"),
    ("2007", "val"),
    ("2007", "test"),
]

# classes = [
#     "aeroplane",
#     "bicycle",
#     "bird",
#     "boat",
#     "bottle",
#     "bus",
#     "car",
#     "cat",
#     "chair",
#     "cow",
#     "diningtable",
#     "dog",
#     "horse",
#     "motorbike",
#     "person",
#     "pottedplant",
#     "sheep",
#     "sofa",
#     "train",
#     "tvmonitor",
# ]


# def convert(size, box):
#     dw = 1.0 / size[0]
#     dh = 1.0 / size[1]
#     x = (box[0] + box[1]) / 2.0
#     y = (box[2] + box[3]) / 2.0
#     w = box[1] - box[0]
#     h = box[3] - box[2]
#     x = x * dw
#     w = w * dw
#     y = y * dh
#     h = h * dh
#     return (x, y, w, h)


# LABELS_ROOT = "VOClabels"


# def convert_annotation(year, image_id):
#     in_file = open("VOCdevkit/VOC%s/Annotations/%s.xml" % (year, image_id))
#     out_file = open("%s/VOC%s/labels/%s.txt" % (LABELS_ROOT, year, image_id), "w")
#     tree = ET.parse(in_file)
#     root = tree.getroot()
#     size = root.find("size")
#     w = int(size.find("width").text)
#     h = int(size.find("height").text)

#     for obj in root.iter("object"):
#         difficult = obj.find("difficult").text
#         cls = obj.find("name").text
#         if cls not in classes or int(difficult) == 1:
#             continue
#         cls_id = classes.index(cls)
#         xmlbox = obj.find("bndbox")
#         b = (
#             float(xmlbox.find("xmin").text),
#             float(xmlbox.find("xmax").text),
#             float(xmlbox.find("ymin").text),
#             float(xmlbox.find("ymax").text),
#         )
#         bb = convert((w, h), b)
#         out_file.write(str(cls_id) + " " + " ".join([str(a) for a in bb]) + "\n")


wd = getcwd()

for year, image_set in sets:
    # if not os.path.exists("%s/VOC%s/labels/" % (LABELS_ROOT, year)):
    #     os.makedirs("%s/VOC%s/labels/" % (LABELS_ROOT, year))
    image_ids = (
        open("VOCdevkit/VOC%s/ImageSets/Main/%s.txt" % (year, image_set))
        .read()
        .strip()
        .split()
    )
    list_file = open("%s_%s.txt" % (year, image_set), "w")
    for image_id in image_ids:
        list_file.write("%s/VOCdevkit/VOC%s/JPEGImages/%s.jpg\n" % (wd, year, image_id))
        # convert_annotation(year, image_id)
    list_file.close()
